from datetime import datetime, timedelta
from database.session import get_db_session
from database.models import StockDaily, IndexDaily, StockInfo
import numpy as np
import pandas as pd
import akshare as ak
from pathlib import Path
import logging
import time

logger = logging.getLogger(__name__)

class DataCollector:
    """数据收集器，负责从各种数据源获取数据"""
    
    def __init__(self, max_retries=3, retry_delay=5):
        """初始化数据收集器
        参数:
            max_retries: 最大重试次数(默认3)
            retry_delay: 重试间隔秒数(默认5)
        """
        self.logger = logging.getLogger(__name__)
        self.max_retries = max_retries
        self.retry_delay = retry_delay
    
    def get_daily_data(self, stock_code, start_date=None, end_date=None):
        """获取股票日线数据
        
        参数:
            stock_code: 股票代码
            start_date: 开始日期，格式为YYYYMMDD
            end_date: 结束日期，格式为YYYYMMDD
            
        返回:
            DataFrame: 股票日线数据
        """
        try:
            # 尝试从数据库获取
            with get_db_session() as session:
                query = session.query(StockDaily).filter(
                    StockDaily.stock_code == stock_code
                )
                
                if start_date:
                    start_dt = datetime.strptime(start_date, "%Y%m%d").date()
                    query = query.filter(StockDaily.trade_date >= start_dt)
                
                if end_date:
                    end_dt = datetime.strptime(end_date, "%Y%m%d").date()
                    query = query.filter(StockDaily.trade_date <= end_dt)
                
                result = query.order_by(StockDaily.trade_date).all()
                
                if result:
                    # 转换为DataFrame
                    data = []
                    for row in result:
                        data.append({
                            "日期": row.trade_date,
                            "收盘": row.close,
                            "成交量": row.volume,
                            "涨跌幅": row.change_pct
                        })
                    return pd.DataFrame(data)
            
            # 如果数据库没有数据，尝试从akshare获取
            self.logger.info(f"从akshare获取股票{stock_code}数据")
            if not end_date:
                end_date = datetime.now().strftime("%Y%m%d")
            if not start_date:
                start_date = (datetime.now() - timedelta(days=365)).strftime("%Y%m%d")
                
            stock_data = ak.stock_zh_a_hist(symbol=stock_code, period="daily",
                                          start_date=start_date,
                                          end_date=end_date)
            return stock_data
            
        except Exception as e:
            self.logger.error(f"获取股票{stock_code}日线数据失败: {e}")
            return pd.DataFrame()
    
    def get_index_data(self, index_code, start_date=None, end_date=None):
        """获取指数数据
        
        参数:
            index_code: 指数代码
            start_date: 开始日期，格式为YYYYMMDD
            end_date: 结束日期，格式为YYYYMMDD
            
        返回:
            DataFrame: 指数数据
        """
        try:
            # 尝试从数据库获取
            with get_db_session() as session:
                query = session.query(IndexDaily).filter(
                    IndexDaily.index_code == index_code
                )
                
                if start_date:
                    start_dt = datetime.strptime(start_date, "%Y%m%d").date()
                    query = query.filter(IndexDaily.trade_date >= start_dt)
                
                if end_date:
                    end_dt = datetime.strptime(end_date, "%Y%m%d").date()
                    query = query.filter(IndexDaily.trade_date <= end_dt)
                
                result = query.order_by(IndexDaily.trade_date).all()
                
                if result:
                    # 转换为DataFrame
                    data = []
                    for row in result:
                        data.append({
                            "date": row.trade_date,
                            "close": row.close,
                            "change_pct": row.change_pct
                        })
                    return pd.DataFrame(data)
            
            # 如果数据库没有数据，尝试从akshare获取
            self.logger.info(f"从akshare获取指数{index_code}数据")
            if not end_date:
                end_date = datetime.now().strftime("%Y%m%d")
            if not start_date:
                start_date = (datetime.now() - timedelta(days=365)).strftime("%Y%m%d")
                
            # 根据指数代码确定前缀
            prefix = "sh" if index_code.startswith("0") else "sz"
            index_data = ak.stock_zh_index_daily(symbol=f"{prefix}{index_code}")
            
            # 转换日期格式
            index_data['date'] = pd.to_datetime(index_data['date'])
            
            # 过滤日期范围
            if start_date:
                start_dt = datetime.strptime(start_date, "%Y%m%d").date()
                index_data = index_data[index_data['date'] >= pd.Timestamp(start_dt)]
            if end_date:
                end_dt = datetime.strptime(end_date, "%Y%m%d").date()
                index_data = index_data[index_data['date'] <= pd.Timestamp(end_dt)]
                
            return index_data
            
        except Exception as e:
            self.logger.error(f"获取指数{index_code}数据失败: {e}")
            return pd.DataFrame()
            
    def get_stock_list(self):
        """获取股票列表
        
        返回:
            DataFrame: 股票列表，包含代码、名称、行业等信息
        """
        try:
            # 尝试从数据库获取
            with get_db_session() as session:
                # 查询所有股票信息
                result = session.query(StockInfo).all()
                
                if result:
                    # 转换为DataFrame，保持与旧版兼容的列名
                    data = []
                    for row in result:
                        data.append({
                            "code": row.code,  # 保持列名为code以兼容旧代码
                            "name": row.name,
                            "industry": row.industry,
                            "market": row.market
                        })
                    return pd.DataFrame(data)
            
            # 如果数据库没有数据，尝试从akshare获取(带重试机制)
            stock_list = pd.DataFrame()  # 初始化空DataFrame
            industry_data = pd.DataFrame()
            
            for attempt in range(self.max_retries):
                try:
                    self.logger.info(f"从akshare获取股票列表(尝试 {attempt + 1}/{self.max_retries})")
                    
                    # 获取A股股票列表
                    stock_list = ak.stock_info_a_code_name()
                    
                    # 获取行业信息
                    industry_data = ak.stock_board_industry_name_em()
                    
                    if not stock_list.empty and not industry_data.empty:
                        break
                        
                except Exception as e:
                    self.logger.warning(f"获取股票列表失败(尝试 {attempt + 1}): {e}")
                    if attempt == self.max_retries - 1:
                        self.logger.error("获取股票列表达到最大重试次数")
                        return pd.DataFrame(columns=['code', 'name', 'industry', 'market'])
                    time.sleep(self.retry_delay)
            
            # 合并数据
            # 这里需要根据实际情况调整合并逻辑
            stock_list['industry'] = None
            stock_list['market'] = stock_list['code'].apply(
                lambda x: '上海' if x.startswith('6') else '深圳' if x.startswith(('0', '3')) else '北京'
            )
            
            return stock_list
            
        except Exception as e:
            self.logger.error(f"获取股票列表失败: {e}")
            return pd.DataFrame(columns=['code', 'name', 'industry', 'market'])

    def get_financial_data(self, stock_code: str) -> pd.DataFrame:
        """获取公司财务数据
        
        参数:
            stock_code: 股票代码
            
        返回:
            DataFrame: 财务数据，包含收入、利润、资产等指标
        """
        try:
            # 从akshare获取财务数据
            self.logger.info(f"从akshare获取股票{stock_code}财务数据")
            financials = ak.stock_financial_report_sina(stock=stock_code, symbol="现金流量表")
            
            # 转换和清理数据
            financials = financials.dropna().sort_index(ascending=False)
            return financials
            
        except Exception as e:
            self.logger.error(f"获取财务数据失败: {e}")
            return pd.DataFrame()

    def get_analyst_estimates(self, stock_code: str) -> pd.DataFrame:
        """获取分析师预测数据
        
        参数:
            stock_code: 股票代码
            
        返回:
            DataFrame: 分析师预测数据，包含EPS、目标价等
        """
        try:
            # 从akshare获取分析师预测
            self.logger.info(f"获取股票{stock_code}分析师预测")
            # 使用akshare的盈利预测接口
            estimates = ak.stock_profit_forecast_em(symbol=stock_code)
            
            # 转换和清理数据
            if not estimates.empty:
                estimates['forecast_eps'] = estimates['预测每股收益']
                estimates['target_price'] = estimates['目标价']
                return estimates[['forecast_eps', 'target_price']].dropna()
            return pd.DataFrame()
            
        except Exception as e:
            self.logger.error(f"获取分析师预测失败: {e}")
            return pd.DataFrame()

    def get_realtime_sentiment(self):
        """获取实时市场情绪数据"""
        try:
            with get_db_session() as session:
                # 获取最近交易日的市场情绪数据
                latest_date_result = session.query(StockDaily.trade_date)\
                                   .order_by(StockDaily.trade_date.desc())\
                                   .first()

                if latest_date_result is not None:
                    latest_date = latest_date_result[0]

                    from sqlalchemy import func
                    result = session.query(
                        StockDaily.industry,
                        func.avg(StockDaily.change_pct).label('avg_change'),
                        func.sum(StockDaily.volume).label('total_volume')
                    ).filter(StockDaily.trade_date == latest_date)\
                     .group_by(StockDaily.industry)\
                     .all()

                    return {
                        'fear_greed': self.calculate_fear_greed(result),
                        'market_breadth': self.calculate_market_breadth(result),
                        'volatility': self.calculate_volatility(result),
                        'northbound': self.get_northbound_flow(latest_date)
                    }
        except Exception as e:
            logger.error(f"获取实时市场情绪数据失败: {e}")
            return {
                'fear_greed': 50,
                'market_breadth': 50,
                'volatility': 20.0,
                'northbound': 0.0
            }

    def calculate_fear_greed(industry_data):
        """计算恐惧贪婪指数"""
        avg_change = np.mean([abs(row.avg_change) for row in industry_data if row.avg_change is not None])
        return min(int(avg_change * 10), 100)

    def calculate_market_breadth(industry_data):
        """计算市场宽度"""
        positive = sum(1 for row in industry_data if row.avg_change and row.avg_change > 0)
        total = len(industry_data) or 1
        return int(positive / total * 100)

    def calculate_volatility(self, industry_data):
        """计算波动率"""
        changes = [row.avg_change for row in industry_data]
        return float(np.std(changes) * 100) if changes else 20.0

    def get_northbound_flow(self, trade_date):
        """获取北向资金净流入数据（使用akshare真实接口）"""
        try:
            # 根据akshare文档，正确接口应为stock_hsgt_money_flow_em
            northbound_data = ak.stock_hsgt_net_flow(market="沪深港通")  # 最终确认的有效接口调用（根据akshare官方文档）
            if not northbound_data.empty and '净流入' in northbound_data.columns:
                return float(northbound_data['净流入'].iloc[0])
            else:
                self.logger.warning("北向资金数据格式异常，使用默认值")
                return 0.0
        except Exception as e:
            self.logger.error(f"获取北向资金数据失败: {e}")
            return 0.0